#include "Material.h"
#include "Debug.h"
#include "RenderThreadStorage.h"
#include "Scene.h"
#include "VectorMath.h"
#include "SobolSequenceGenerator.h"
#include "GI.h"
#include "Windows/GxThreads.h"
#include "FastRandom.h"
#include <assert.h>
#include <ppl.h>

using namespace GxLibMath;
using namespace Concurrency;

namespace RayTracePro
{
	const float ConeFilterConstant = 1.1f;

	const unsigned int BitMasks[32] = {1, 1<<1, 1<<2, 1<<3, 1<<4, 1<<5, 1<<6, 1<<7, 1<<8, 1<<9,
									  1<<10, 1<<11, 1<<12, 1<<13, 1<<14, 1<<15, 1<<16, 1<<17, 1<<18, 1<<19,
									  1<<20, 1<<21, 1<<22, 1<<23, 1<<24, 1<<25, 1<<26, 1<<27, 1<<28, 1<<29,
									  1<<30, 1<<31};
	
	// implemented in StandardShaders.h
	bool EvalLightRadiance(Vec3 & rs, Vec3 & l, RayTraceSystem * system, Scene * scene, Light * light, const Vec3 & p);
	
	__forceinline void DiffuseReflectPhoton(float & cosTheta, Vec3 & dirOut, Vec3 & posOut, SobolGenerator & rand, Vec3 & position, Vec3 & normal, float shadowBias)
	{
		float phi = (float)(2 * PI * rand.NextX());
		Vec3 nDir;
		float val = (float)rand.NextY();
		nDir.y = sqrt(val);
		cosTheta = nDir.y;
		float s = sqrt(1-val);
		nDir.x = cos(phi) * s;
		nDir.z = sin(phi) * s;
		Vec3 tangent, bitangent;
		GetOrthoVec(tangent, normal);
		Vec3::Cross(bitangent, normal, tangent);
		dirOut.x = tangent.x * nDir.x + normal.x * nDir.y + bitangent.x * nDir.z;
		dirOut.y = tangent.y * nDir.x + normal.y * nDir.y + bitangent.y * nDir.z;
		dirOut.z = tangent.z * nDir.x + normal.z * nDir.y + bitangent.z * nDir.z;
		Vec3 tmp;
		Vec3::Scale(tmp, normal, shadowBias);
		posOut = position;
		posOut += tmp;
	}
	__forceinline void DiffuseReflectPhoton(float & cosTheta, Vec3 & dirOut, Vec3 & posOut, FastRandom & randGen, Vec3 & position, Vec3 & normal, float shadowBias)
	{
		float phi = 2 * PI * randGen.NextFloat();
		Vec3 nDir;
		float val = randGen.NextFloat();
		nDir.y = sqrt(val);
		cosTheta = nDir.y;
		float s = 1-val;
		nDir.x = cos(phi) * s;
		nDir.z = sin(phi) * s;
		Vec3 tangent, bitangent;
		GetOrthoVec(tangent, normal);
		Vec3::Cross(bitangent, normal, tangent);
		dirOut.x = tangent.x * nDir.x + normal.x * nDir.y + bitangent.x * nDir.z;
		dirOut.y = tangent.y * nDir.x + normal.y * nDir.y + bitangent.y * nDir.z;
		dirOut.z = tangent.z * nDir.x + normal.z * nDir.y + bitangent.z * nDir.z;
		Vec3 tmp;
		Vec3::Scale(tmp, normal, shadowBias);
		posOut = position;
		posOut += tmp;
	}
	
	template<typename T>
	void ComputeBBox(BBox & bbox, Array<T, true> & Photons)
	{
		bbox.Init();
		for (int i = 1; i<Photons.Count(); i++)
		{
			bbox.Union(Photons[i].Position);
		}
	}

	template<typename T>
	void BalanceGenericPhotons(Array<T, true> & Photons, BBox & bbox)
	{
		int storedPhotons = Photons.Count()-1;
		if (storedPhotons>1)
		{
			// allocate two temporary arrays for the balancing procedure
			Array<T *, true> pa1;
			Array<T *, true> pa2;
			pa1.SetSize(storedPhotons + 1);
			pa2.SetSize(storedPhotons + 1);
			
			int i;
			for (i=0; i<=storedPhotons; i++)
				pa2[i] = &Photons[i];

			BalanceSegment(bbox, pa1, pa2, 1, 1, storedPhotons);
			pa2.Clear(); pa2.Compress();
			// reorganize balanced kd-tree (make a heap)
			int d, j=1, foo=1;
			Photon foo_photon = Photons[j];
			for (i=1; i<=storedPhotons; i++)
			{
				d=pa1[j]-Photons.Buffer();
				pa1[j] = NULL;
				if (d != foo)
					Photons[j] = Photons[d];
				else
				{
					Photons[j] = foo_photon;
					if (i<storedPhotons)
					{
						for (;foo<=storedPhotons; foo++)
							if (pa1[foo] != NULL)
								break;
						foo_photon = Photons[foo];
						j = foo;
					}
					continue;
				}
				j = d;
			}
		}
	}

	void PhotonMap::Balance()
	{	
		ComputeBBox(bbox, Photons);
		BalanceGenericPhotons(Photons, bbox);
		halfStoredPhotons = (Photons.Count()-2)/2;
	}

	template<typename T>
	void BalanceSegment(BBox & bbox, Array<T *, true> & pbal, Array<T *, true> & porg, const int index, const int start, const int end )
	{
		// compute new median
		int median=1;
		while ((4*median) <= (end-start+1))
			median += median;
		if ((3*median) <= (end-start+1))
		{
			median += median;
			median += start-1;
		}
		else		
			median = end-median+1;

		// find axis to split along
		int axis=2;
		if ((bbox.Max[0]-bbox.Min[0])>(bbox.Max[1]-bbox.Min[1]) &&
			(bbox.Max[0]-bbox.Min[0])>(bbox.Max[2]-bbox.Min[2]))
			axis=0;
		else if ((bbox.Max[1]-bbox.Min[1])>(bbox.Max[2]-bbox.Min[2]))
			axis=1;

		// partition photon block around the median
		MedianSplit(porg, start, end, median, axis);
		pbal[index] = porg[median];
		pbal[index]->SetSplitDimension(axis);

		// recursively balance the left and right block
		if ( median > start )
		{
			// balance left segment
			if ( start < median-1 )
			{
				const float tmp = bbox.Max[axis];
				bbox.Max[axis] = pbal[index]->Position[axis];
				BalanceSegment( bbox, pbal, porg, 2*index, start, median-1 );
				bbox.Max[axis] = tmp;
			}
			else
			{
				pbal[2*index] = porg[start];
			}
		}
		if (median < end)
		{
			// balance right segment
			if (median+1 < end)
			{
				const float tmp = bbox.Min[axis];		 
				bbox.Min[axis] = pbal[index]->Position[axis];
				BalanceSegment( bbox, pbal, porg, 2*index+1, median+1, end );
				bbox.Min[axis] = tmp;
			}
			else
			{
				pbal[2*index+1] = porg[end];
			}
		} 
	}

	template<typename T>
	void MedianSplit(Array<T *, true> &p,
		const int start,               // start of photon block in array
		const int end,                 // end of photon block in array
		const int median,              // desired median number
		const int axis )               // axis to split along
	{
#define SWAP(ph,a,b) { T *ph2=ph[a]; ph[a]=ph[b]; ph[b]=ph2; }		
		int left = start;
		int right = end;
		while (right > left)
		{
			const float v = p[right]->Position[axis];
			int i=left-1;
			int j=right;
			while (true)
			{
				while (p[++i]->Position[axis] < v);
				while (p[--j]->Position[axis] > v && j>left);
				if (i >= j)
					break;
				SWAP(p,i,j);
			}
			SWAP(p,i,right);
			if (i >= median)
				right=i-1;
			if (i <= median)
				left=i+1;
		}
	}

	void PhotonMap::SearchPhotons(PhotonSearchResult & rs, const Vec3 & position, float radius)
	{
		int searchStack[64];
		int stackPtr = 0;
		float radius2 = radius * radius;
		int index = 1;
		while(true)
		{
			Photon *p = &Photons[index];
			// compute squared distance between current photon and np->pos
			Vec3 distVec;
			Vec3::Subtract(distVec, p->Position, position);
			float dist2 = Vec3::Dot(distVec, distVec);

			if (dist2 < rs.Distances[0])
			{
				if (rs.Photons.Count() < rs.Photons.GetCapacity())
				{
					// heap is not full; use array
					rs.Photons.Add(p);
					rs.Distances.Add(dist2);
				}
				else
				{
					int j,parent;
					int totalPhotons = rs.Photons.Count()-1;
					if (!rs.HeapCreated)
					{
						float dst2;
						Photon *photon;
						int halfPhotons = totalPhotons>>1;
						for (int k=halfPhotons; k>=1; k--)
						{
							parent=k;
							photon = rs.Photons[k];
							dst2 = rs.Distances[k];
							while (parent <= halfPhotons)
							{
								j = parent+parent;
								if (j<totalPhotons && rs.Distances[j]<rs.Distances[j+1])
									j++;
								if (dst2>=rs.Distances[j])
									break;
								rs.Distances[parent] = rs.Distances[j];
								rs.Photons[parent] = rs.Photons[j];
								parent=j;
							}
							rs.Distances[parent] = dst2;
							rs.Photons[parent] = photon;
						}
						rs.HeapCreated = 1;
					}
					// insert new photon into max heap
					// delete largest element, insert new and reorder the heap
					parent=1;
					j = 2;
					while (j <= totalPhotons)
					{
						if (j < totalPhotons && rs.Distances[j] < rs.Distances[j+1] )
							j++;
						if ( dist2 > rs.Distances[j] )
							break;
						rs.Distances[parent] = rs.Distances[j];
						rs.Photons[parent] = rs.Photons[j];
						parent = j;
						j += j;
					}
					rs.Photons[parent] = p;
					rs.Distances[parent] = dist2;
					rs.Distances[0] = rs.Distances[1];
				}
			}
			float dist1;
			if (index <= halfStoredPhotons)
			{
				dist1 = position[p->GetSplitDimension()] - p->Position[p->GetSplitDimension()];
				if (dist1>0.0)
				{ 
					if (dist1 < radius)
					{
						searchStack[stackPtr++] = 2*index;
					}
					// if dist1 is positive search right plane
					index = 2*index+1;
				}
				else
				{         
					if (dist1 < radius)
						searchStack[stackPtr++] = 2*index+1;
					// dist1 is negative search left first
					index = 2*index;
				}
			}
			else
			{
				if (stackPtr == 0)
					break;
				index = searchStack[--stackPtr];
			}
		}
		
	}

	void PhotonMap::PrecomputeIrradiance(Array<Importon, true>& importons, Scene * scene, int photonsPerEstimate, float radius)
	{
		Irradiances.SetCapacity(Photons.Count()<<2);
		Irradiances.Add(IrradiancePhoton());
		int threadCount = DetectCpuCount();
		int stepSize = importons.Count()/threadCount;
		Array<RefPtr<PhotonSearchResult>> buffers;
		Array<RefPtr<Array<IrradiancePhoton,true>>> photonBuffers;
		buffers.SetSize(threadCount);
		photonBuffers.SetSize(threadCount);
		for (int i = 0; i<threadCount; i++)
		{
			buffers[i] = new PhotonSearchResult(photonsPerEstimate);
			photonBuffers[i] = new Array<IrradiancePhoton, true>();
			photonBuffers[i]->SetCapacity(stepSize<<2);
		}

		/*importons.Clear();
		for (int i = 1; i<Photons.Count(); i+=4)
		{
			Vec3 normal;
			Photons[i].GetNormal(normal, dirTable);
			Importon imp;
			imp.Position = Photons[i].Position;
			imp.Normal = normal;
			importons.Add(imp);
		}*/
		
		parallel_for (0, threadCount, [&](int c)
		{
			PhotonSearchResult & buffer = *buffers[c];
			Array<IrradiancePhoton,true> & photonBuffer = *photonBuffers[c];
			for (int i = c*stepSize; i<min((c+1)*stepSize, importons.Count()); i++)
			{
				buffer.Distances.Clear();
				buffer.Photons.Clear();
				buffer.Distances.Add(radius*radius);
				buffer.Photons.Add(0);
				buffer.HeapCreated = 0;
				IrradiancePhoton ip;
				ip.Position = importons[i].Position;
				ip.SetNormal(importons[i].Normal);
				Vec3 rs;
				EstimateIrradianceOriginal(rs, buffer, importons[i].Position, importons[i].Normal, radius);
				Light ** lights;
				int lightCount;
				scene->GetLights(ip.Position, lights, lightCount);
				for (int i = 0; i < lightCount; i++)
				{
					Vec3 rad, l;
					Vec3 pos, tmp;
					Vec3::Scale(tmp, importons[i].Normal, scene->GetSystem()->ShadowBias);
					Vec3::Add(pos, ip.Position, tmp);
					if (EvalLightRadiance(rad, l, scene->GetSystem(), scene, lights[i], pos))
					{
						float lambert = Max(0.0f, Vec3::Dot(l, importons[i].Normal));
						Vec3::Scale(rad, rad, lambert);
					}
					rs.x += rad.x; rs.y += rad.y; rs.z += rad.z;
				}
				
				ip.SetPower(rs);
				photonBuffer.Add(ip);
			}
		});
		for (int i = 0; i<threadCount; i++)
		{
			Irradiances.Add(*photonBuffers[i]);
		}
		if (Irradiances.Count()>1)
		{
			BBox bbox;
			ComputeBBox(bbox, Irradiances);
			BalanceGenericPhotons(Irradiances, bbox);
			halfIrradiancePhotons = (Irradiances.Count()-2)/2;
		}
	}

	void PhotonMap::EstimateIrradiance(Vec3 & rs, PhotonSearchResult & buffer, const Vec3 & position, const Vec3 & normal, float radius)
	{
		rs.SetZero();
		if (Irradiances.Count() <= 1)
			return;
		// Search neareast photon
		int searchStack[64];
		int stackPtr = 0;
		float radius2 = radius * radius;
		int index = 1;
		float minDist = radius2;
		bool found = false;
		while(true)
		{
			IrradiancePhoton *p = &Irradiances[index];
			// compute squared distance between current photon and np->pos
			Vec3 distVec;
			Vec3::Subtract(distVec, p->Position, position);
			float dist2 = Vec3::Dot(distVec, distVec);

			if (dist2 < minDist)
			{
				Vec3 n;
				p->GetNormal(n, dirTable);
				if (Vec3::Dot(normal, n) > 0.7f)
				{
					// find a good irradiance photon, return irradiance
					Vec4 power;
					p->GetPower(power);
					power.xyz(rs);
					minDist = dist2;
					found = true;
				}
			}
			float dist1;
			if (index <= halfIrradiancePhotons)
			{
				dist1 = position[p->GetSplitDimension()] - p->Position[p->GetSplitDimension()];
				if (dist1>0.0)
				{ 
					if (dist1 < radius)
					{
						searchStack[stackPtr++] = 2*index;
					}
					// if dist1 is positive search right plane
					index = 2*index+1;
				}
				else
				{         
					if (dist1 < radius)
						searchStack[stackPtr++] = 2*index+1;
					// dist1 is negative search left first
					index = 2*index;
				}
			}
			else
			{
				if (found || stackPtr == 0)
					break;
				index = searchStack[--stackPtr];
			}
		}
	}

	void PhotonMap::EstimateIrradianceOriginal(Vec3 & rs, PhotonSearchResult & buffer, const Vec3 & pos, const Vec3 & normal, float maxRadius)
	{
		rs.SetZero();

		if (Photons.Count() <= 10)
			return;

		buffer.HeapCreated = 0;
		buffer.Distances.Clear(); 
		buffer.Photons.Clear();
		buffer.Distances.Add(maxRadius * maxRadius);
		buffer.Photons.Add(0);
		// locate the nearest photons
		SearchPhotons(buffer, pos, maxRadius);
		// if less than 8 photons return
		/*if (buffer.Photons.Count()<8)
		{
			return;
		}*/

		Vec3 pdir;
		int numPhotons = buffer.Photons.Count()-1;
		//float dist = sqrt(buffer.Distances[0]);
		//float invKr = 1.0f/(ConeFilterConstant * dist);
		// sum irradiance from all photons
		for (int i=1; i<=numPhotons; i++)
		{
			Photon *p = buffer.Photons[i];
			// the photon_dir call and following if can be omitted (for speed)
			// if the scene does not have any thin surfaces
			p->GetDirection(pdir, dirTable);
			if ((pdir[0]*normal[0]+pdir[1]*normal[1]+pdir[2]*normal[2]) < 0.0f )
			{
				Vec4 power;
				p->GetPower(power);
				//Vec3 distVec;
				//Vec3::Subtract(distVec, pos, p->Position);
				//float weight = 1.0f - (distVec.Length()*invKr);
				rs[0] += power[0];//*weight;
				rs[1] += power[1];//*weight;
				rs[2] += power[2];//*weight;
			}
		}

		const float tmp = (1.0f/PI)/buffer.Distances[0];//(1.0f/(PI*(1.0f-2.0f/(3.0f*ConeFilterConstant))))/(buffer.Distances[0]);
		rs[0] *= tmp;
		rs[1] *= tmp;
		rs[2] *= tmp;
	}

	class PhotonMapEvaluator
	{
	private:
		Scene * scene;
		RayTraceSystem * system;
		PhotonMap * photonCollection;
		struct LightSummary
		{
			Vec3 Center;
			Vec3 ScaledPower;
			Vec3 Direction1, Direction2;
			float Size, InvSize;
			float Weight;
			int Photons;
			unsigned int ProjectionMap[32];
			void ProjectSphereDirectionalLight(const Vec4 & sphere)
			{
				
				Vec3 pos;
				sphere.xyz(pos);
				Vec3 d;
				Vec3::Subtract(d, pos, Center);
				float x = Vec3::Dot(d, Direction1)*InvSize;
				float y = Vec3::Dot(d, Direction2)*InvSize;
				float r = sphere.w * InvSize;
				int minX = min(31, max(0, (int)floor(32*(x-r))));
				int maxX = min(31, max(0, (int)ceil(32*(x+r))));
				int minY = min(31, max(0, (int)floor(32*(y-r))));
				int maxY = min(31, max(0, (int)ceil(32*(y+r))));
				for (int i = minX; i<=maxX; i++)
				{
					unsigned int bit = ProjectionMap[i];
					for (int j = minY; j<=maxY; j++)
					{
						bit |= BitMasks[j];
					}
					ProjectionMap[i] = bit;
				}
			}
			void ProjectSpherePointLight(const Vec4 & sphere)
			{
				// we are not using projection maps for point lights.
				// projections maps are unlikely to have any value, as it's often used 
				// in closed scenes.
			}
		};
		Array<LightSummary, true> lightSummaries;
		struct CreateSceneBoundingSpheresLeafFunction
		{
			int bufferSize;
			Array<Vec4, true> & buffer;
			KdTree<Geometry*> *tree;
			CreateSceneBoundingSpheresLeafFunction(int _bufferSize, Array<Vec4, true> & _buffer)
				:bufferSize(_bufferSize), buffer(_buffer)
			{}
			void operator()(BBox & curBox, int nodeId, int depth)
			{
				Vec4 sphere;
				if (depth == 0)
				{
					Vec3 position, range;
					Vec3::Add(position, curBox.Min, curBox.Max);
					Vec3::Scale(position, position, 0.5f);
					sphere.x = position.x; sphere.y = position.y; sphere.z = position.z;
					Vec3::Subtract(range, curBox.Max, curBox.Min);
					sphere.w = range.Length() * 0.5f;
					buffer.Add(sphere);
				}
				else
				{
					depth -= Log2Ceil(tree->Nodes[nodeId].GetElementCount());
					if (depth <= 0 || tree->Nodes[nodeId].GetElementCount() > 3)
					{
						Vec3 position, range;
						Vec3::Add(position, curBox.Min, curBox.Max);
						Vec3::Scale(position, position, 0.5f);
						sphere.x = position.x; sphere.y = position.y; sphere.z = position.z;
						Vec3::Subtract(range, curBox.Max, curBox.Min);
						sphere.w = range.Length() * 0.5f;
						buffer.Add(sphere);
					}
					else
					{
						if (tree->Nodes[nodeId].GetElementCount() == 1)
						{
							tree->Nodes[nodeId].OneElement->Element->ComputeBoundingSpheres(depth, buffer);
						}
						else
						{
							for (int i = 0; i<(int)tree->Nodes[nodeId].GetElementCount(); i++)
							{
								tree->Nodes[nodeId].Elements[i]->Element->ComputeBoundingSpheres(depth, buffer);
							}
						}
					}
				}
			}
		};
		void SummaryLights(int numPhotons)
		{
			// Create simplified scene bounding spheres;
			int spheresCount = 0;
			Array<Vec4, true> spheres;
			spheres.SetCapacity(32);
			CreateSceneBoundingSpheresLeafFunction cleaf(spheres.Count(), spheres);
			cleaf.tree = &scene->kdTree;
			ComputeKdTreeBoundingSpheres(scene->kdTree, 5, spheres, cleaf);
			// Step 1, gather light properties
			lightSummaries.SetSize(scene->lights.Count());
			memset(lightSummaries.Buffer(), 0, sizeof(LightSummary)*lightSummaries.Count());
			Vec3 tmp;
			Vec3::Subtract(tmp, scene->kdTree.Bounds.Max, scene->kdTree.Bounds.Min);
			float sceneRadius = tmp.Length() * 0.5f;
			float sumWeight = 0.0f;
			for (int i = 0; i < scene->lights.Count(); i++)
			{
				Light & light = *scene->lights[i];
				LightSummary &sum = lightSummaries[i];
				if (light.IsDirectional)
				{
					// Direction1 and Direction2 are two orthogonal axis that forms the photon emission plane
					GetOrthoVec(sum.Direction1, light.Direction);
					Vec3::Normalize(sum.Direction1, sum.Direction1);
					Vec3::Cross(sum.Direction2, sum.Direction1, light.Direction);
					Vec3::Normalize(sum.Direction2, sum.Direction2);
					sum.Size = sceneRadius;
					// A directional light stores power/square meters in its Intensity field
					// Therefore the power of a directional light is Intensity multiplied by projected area
					Vec3::Scale(sum.ScaledPower, light.Intensity, sceneRadius*sceneRadius*4);
					Vec3::Add(sum.Center, scene->kdTree.Bounds.Max, scene->kdTree.Bounds.Min);
					Vec3::Scale(sum.Center, sum.Center, 0.5f);
					Vec3 tmp;
					Vec3::Scale(tmp, light.Direction, sum.Size);
					sum.Center -= tmp;
					Vec3::Scale(tmp, sum.Direction1, -sum.Size);
					sum.Center += tmp;
					Vec3::Scale(tmp, sum.Direction2, -sum.Size);
					sum.Center += tmp;
					sum.Size *= 2.0f;
					sum.InvSize = 1.0f/sum.Size;
				}
				else
				{
					sum.Direction1 = Vec3(1.0f, 0.0f, 0.0f);
					sum.Direction2 = Vec3(0.0f, 1.0f, 0.0f);
					sum.Center = light.Position;
					Vec3::Scale(sum.ScaledPower, light.Intensity, 4.0f * PI);
					sum.Size = 0.0f;
				}
				sum.Weight = (sum.ScaledPower.x + sum.ScaledPower.y + sum.ScaledPower.z)*0.33333f;
				
			}
			// Step 2: create projection maps
			for (int i = 0; i<spheres.Count(); i++)
			{
				for (int j = 0; j<lightSummaries.Count(); j++)
				{
					if (scene->lights[j]->IsDirectional)
						lightSummaries[j].ProjectSphereDirectionalLight(spheres[i]); // not thread safe
					else
						lightSummaries[j].ProjectSpherePointLight(spheres[i]); // not thread safe
				}
			}
			// Step 3: Scale Power based projecton map
			for (int i = 0; i<lightSummaries.Count(); i++)
			{
				if (scene->lights[i]->IsDirectional)
				{
					int activeCells = 0;
					for (int j = 0; j<32; j++)
					{
						activeCells += ones32(lightSummaries[i].ProjectionMap[j]);
					}
					float powerScale = activeCells / (float)(32*32);
					Vec3::Scale(lightSummaries[i].ScaledPower, lightSummaries[i].ScaledPower, powerScale);
					lightSummaries[i].Weight *= powerScale;
					
				}
				sumWeight += lightSummaries[i].Weight;
			}
			// Step 4: assign photon headcounts for each light
			double invSumWeight = 1.0 / sumWeight;
			for (int i = 0; i<scene->lights.Count(); i++)
			{
				lightSummaries[i].Photons = Ceil(numPhotons * lightSummaries[i].Weight * invSumWeight);
			}
		}

		struct CellId
		{
			short X, Y;
		};

		void EmitPhotons(int numPhotons, int maxDepth)
		{
			SummaryLights(numPhotons);
			Array<CellId, true> activeCells;
			for (int i = 0; i<scene->lights.Count(); i++)
			{
				activeCells.Clear();
				Light & light = *scene->lights[i];
				if (light.IsDirectional)
				{
					for (int j = 0; j<32; j++)
					{
						for (int k = 0; k<32; k++)
						{
							if (lightSummaries[i].ProjectionMap[j] & BitMasks[k])
							{
								CellId activeCell;
								activeCell.X = j; activeCell.Y = k;
								activeCells.Add(activeCell);
							}
						}
					}
					int photonsPerCell = lightSummaries[i].Photons / activeCells.Count();
					int totalPhotons = photonsPerCell * activeCells.Count();
					float cellSize = lightSummaries[i].Size / 32.0f;
					Vec3 photonPower;
					Vec3::Scale(photonPower, lightSummaries[i].ScaledPower, 1.0f/totalPhotons);
					parallel_for(0, activeCells.Count(), [&](int c)
					{
						FastRandom randGen(c*39287);
						PhotonMap::PhotonCollectionBuffer buffer;
						photonCollection->CreatePhotonCollectionBuffer(buffer, 1024);
						//SobolGenerator random;
					
						CellId cell = activeCells[c];
						for (int p = 0; p<photonsPerCell; p++)
						{
							float x = cellSize * (float)randGen.NextFloat();
							float y = cellSize * (float)randGen.NextFloat();
							x += cellSize * cell.X;
							y += cellSize * cell.Y;
							Vec3 pos, tmp;
							pos = lightSummaries[i].Center;
							Vec3::Scale(tmp, lightSummaries[i].Direction1, x);
							pos += tmp, 
							Vec3::Scale(tmp, lightSummaries[i].Direction2, y);
							pos += tmp;
							TracePhoton(buffer, pos, scene->lightPtrs[i]->Direction, photonPower, maxDepth, randGen);
						}
						photonCollection->AddBuffer(buffer);
					}
					);
				}
				else
				{
					int layers = DetectCpuCount();
					if (light.Cutoff)
						throw "Photon emitting from spot light is not implemented";
					float layerSize = 2.0f/layers;
					int photonsPerLayer = lightSummaries[i].Photons / layers;
					int totalPhotons = photonsPerLayer * layers;
					Vec3 photonPower;
					Vec3::Scale(photonPower, lightSummaries[i].ScaledPower, 1.0f/totalPhotons);
					parallel_for(0, layers, 1, [&](int lid)
					{
						SobolGenerator random;
						FastRandom randGen(lid*39287);
						// uniformly sampling the sphere
						/*
						  The trig method.  This method works only in 3-space, but it is
						  very fast.  It depends on the slightly counterintuitive fact (see
						  proof below) that each of the three coordinates of a uniformly
						  distributed point on S^2 is uniformly distributed on [-1,1] (but
						  the three are not independent, obviously).  Therefore, it
						  suffices to choose one axis (Z, say) and generate a uniformly
						  distributed value on that axis.  This constrains the chosen point
						  to lie on a circle parallel to the X-Y plane, and the obvious
						  trig method may be used to obtain the remaining coordinates.

						(a) Choose z uniformly distributed in [-1,1].
						(b) Choose t uniformly distributed on [0, 2*pi).
						(c) Let r = sqrt(1-z^2).
						(d) Let x = r * cos(t).
						(e) Let y = r * sin(t).
						*/
						PhotonMap::PhotonCollectionBuffer buffer;
						photonCollection->CreatePhotonCollectionBuffer(buffer, 1024);
						float beginZ = -1.0f + lid*layerSize;
						for (int p = 0; p<photonsPerLayer; p++)
						{
							float z = beginZ + layerSize * (float)random.NextX();
							float r = sqrt(1.0f - z*z);
							float t = (float)random.NextY() * 2.0f * PI;
							float x = r * cos(t);
							float y = r * sin(t);
							Vec3 dir(x, y, z);
							TracePhoton(buffer, light.Position, dir, photonPower, maxDepth, randGen);
						}
						photonCollection->AddBuffer(buffer);
					});
				}
			}
		}

		void TracePhoton(PhotonMap::PhotonCollectionBuffer & buffer, const Vec3 & pos, const Vec3 & dir, const Vec3 & power, int depth, FastRandom & randGen)
		{
			DifferentialGeometry dgeom;
			Ray ray;
			RayDifferential rayDiff;
			float t;
			ray.Direction = dir;
			ray.Origin = pos;
			ray.tMin = Epsilon;
			ray.tMax = FLT_MAX;
			ray.ReciprocalDirection.x = 1.0f/dir.x;
			ray.ReciprocalDirection.y = 1.0f/dir.y;
			ray.ReciprocalDirection.z = 1.0f/dir.z;
			memset(&rayDiff, 0, sizeof(RayDifferential));
			if (scene->TraceRayCore(dgeom, ray, rayDiff, t))
			{
				if (dgeom.Material
					&& dgeom.Material->GI.Lambert)
				{
					Vec4 diffuse;
					{
						auto & giMat = dgeom.Material->GI;
						giMat.Diffuse.GetSimplified(&diffuse, &dgeom, &ray, &rayDiff, scene, depth);
						if (scene->GetSystem()->PhotonTraceDepth > depth)
						{
							Photon photon;
							photon.Flags = 0;
							photon.Position = dgeom.Position;
							photon.SetDirection(ray.Direction);
							photon.SetPower(power);
							photon.SetNormal(dgeom.Normal);
							buffer.Add(photon);
						}
					}
					float prob = (diffuse.x + diffuse.y + diffuse.z) * 0.3333333f;
					float invProb = 1.0f/prob;
					if (randGen.NextFloat()<=prob && depth)
					{
						// reflect
						Vec3 nPower;
						nPower.x = power.x*diffuse.x*invProb;
						nPower.y = power.y*diffuse.y*invProb;
						nPower.z = power.z*diffuse.z*invProb;
						Vec3 photonDir;
						Vec3 pos;
						float cosTheta;
						DiffuseReflectPhoton(cosTheta, photonDir, pos, randGen, dgeom.Position, dgeom.Normal, system->ShadowBias);
						TracePhoton(buffer, pos, photonDir, nPower, depth-1, randGen);
					}
				}
			}
		}
	public:
		PhotonMapEvaluator(RayTraceSystem * system, Scene * scene)
		{
			this->scene = scene;
			this->system = system;
		}
		
		void CreatePhotonMap(int numPhotons, int maxDepth)
		{
			photonCollection = new PhotonMap(numPhotons);
			DebugTimming::Start("Shooting photons...");
			EmitPhotons(numPhotons, maxDepth);
			DebugTimming::End();
			printf("%d photons stored.\n", photonCollection->Photons.Count());
			DebugTimming::Start("Balancing photon map...");
			photonCollection->Balance();
			DebugTimming::End();
			
		}

		PhotonMap * GetPhotonMap()
		{
			return photonCollection;
		}
	};

	
	void TraceImporton(Scene * scene, Array<Importon, true> & importons, Ray & ray, RayDifferential & rayDiff, FastRandom & randGen, float shadowBias, int depth)
	{
		DifferentialGeometry dgeom;
		float t;
		if (scene->TraceRayCore(dgeom, ray, rayDiff, t))
		{
			if (dgeom.Material
				&& dgeom.Material->GI.Lambert)
			{
				Importon imp;
				imp.Position = dgeom.Position;
				imp.Normal = dgeom.Normal;
				importons.Add(imp);
				if (depth > 0)
				{
					for (int i = 0; i<16; i++)
					{
						Vec3 photonDir, pos;
						float cosTheta;
						DiffuseReflectPhoton(cosTheta, photonDir, pos, randGen, dgeom.Position, dgeom.Normal, shadowBias);
					
						Ray ray;
						ray.Direction = photonDir;
						ray.Origin = pos;
						ray.tMin = Epsilon;
						ray.tMax = FLT_MAX;
						ray.ReciprocalDirection.x = 1.0f/ray.Direction.x;
						ray.ReciprocalDirection.y = 1.0f/ray.Direction.y;
						ray.ReciprocalDirection.z = 1.0f/ray.Direction.z;
						TraceImporton(scene, importons, ray, rayDiff, randGen, shadowBias, depth-1);
					}
				}
			}
		}
	}

	void EmitImportons(Scene * scene, Array<Importon, true> & importons, CameraProjection & proj)
	{
		int w = proj.GetScreenWidth();
		int h = proj.GetScreenHeight();
		int taskW = w/256;
		int taskH = h/256;
		int tasks = taskW * taskH;
		GxWinSystem::Threading::SpinLock lock;
		int shootStep = proj.GetScreenWidth()/scene->GetSystem()->ImportonResolution;
		importons.SetCapacity(scene->GetSystem()->ImportonResolution*scene->GetSystem()->ImportonResolution);
		parallel_for (0, tasks, [&](int c)
		{
			div_t dv = div(c, taskW);
			int ty = dv.quot;
			int tx = dv.rem;
			Array<Importon, true> importonBuffer;
			importonBuffer.SetCapacity(1024);
			FastRandom randGen((c+1) * 8623173);
			for (int i = tx * 256; i<min((tx+1)*256, w); i+=shootStep)
			{
				for (int j = ty * 256; j<min((ty+1)*256, w); j+=shootStep)
				{
					Ray r;
					RayDifferential rayDiff;
					proj.GenerateRay(r, rayDiff, (float)i, (float)j);
					TraceImporton(scene, importonBuffer, r, rayDiff, randGen, scene->GetSystem()->ShadowBias, 1);
				}
			}
			lock.Lock();
			importons.Add(importonBuffer);
			lock.Unlock();
		});
	}

	void Scene::PrecomputeGI()
	{
		Vec3 size;
		Vec3::Subtract(size, this->kdTree.Bounds.Max, this->kdTree.Bounds.Min);
		float sizef = size.Length();
		this->globalPhotonMapRadius = min(sizef/50.0f, 10000 * sizef / system->TotalPhotons);
		this->globalPhotonMapRadius = this->GlobalPhotonMapRadianceModifier.GetRadius(this->globalPhotonMapRadius);
		PhotonMapEvaluator eval(system, this);
		eval.CreatePhotonMap(system->TotalPhotons, system->PhotonTraceDepth);
		this->globalPhotonMap = eval.GetPhotonMap();
	}

	union Int2Float
	{
		float FloatVal;
		int IntVal;
	};
	void GIMaterial::EvaluateGI(Vec4 & result, DifferentialGeometry * dgeom, const Ray * ray, const RayDifferential * rayDiff, Scene * scene, int depth)
	{
		Int2Float cvt;
		cvt.FloatVal = dgeom->Position.x + dgeom->Position.y;
		int xseed = cvt.IntVal;
		cvt.FloatVal = dgeom->Position.y + dgeom->Position.z;
		int yseed = cvt.IntVal;
		SobolGenerator rand((xseed<<16)^yseed, (yseed<<14)^xseed);
		float shadowBias = scene->GetSystem()->ShadowBias;
		float sumWeight = 0.0f;
		result.SetZero();
		result.w = 1.0f;
		for (int i = 0; i<scene->GetSystem()->FinalGatherRays; i++)
		{
			Ray r;
			float cosTheta;
			DiffuseReflectPhoton(cosTheta, r.Direction, r.Origin, rand, dgeom->Position, dgeom->Normal, shadowBias);
			r.tMin = Epsilon;
			r.tMax = FLT_MAX;
			r.ReciprocalDirection.x = 1.0f/r.Direction.x;
			r.ReciprocalDirection.y = 1.0f/r.Direction.y;
			r.ReciprocalDirection.z = 1.0f/r.Direction.z;
			DifferentialGeometry inter;
			float t;
			if (scene->TraceRayCore(inter, r, *rayDiff, t))
			{
				if (inter.Material->GI.Lambert)
				{
					Vec4 rs, diffuse;
					inter.Material->GI.EstimateGI(rs, &inter, &r, rayDiff, scene, depth);
					inter.Material->GI.Diffuse.GetSimplified(&diffuse, &inter, &r, rayDiff, scene, depth);
					//sumWeight += cosTheta;
					result.x += rs.x*diffuse.x;//*cosTheta;
					result.y += rs.y*diffuse.y;//*cosTheta;
					result.z += rs.z*diffuse.z;//*cosTheta;
				}
			}
			else if (scene->Environment)
			{
				//sumWeight += cosTheta;
				Vec4 env;
				scene->Environment->Shade(env, r.Direction, Vec3(0.1f), Vec3(0.1f));
				result.x += env.x;//*cosTheta;
				result.y += env.y;//*cosTheta;
				result.z += env.z;//*cosTheta;
			}
		}
		float invSumWeight = 1.0f/scene->GetSystem()->FinalGatherRays;
		result.x *= invSumWeight;
		result.y *= invSumWeight;
		result.z *= invSumWeight;
	}

	void GIMaterial::EstimateGI(Vec4 & result, DifferentialGeometry * dgeom, const Ray * ray, const RayDifferential * rayDiff, Scene * scene, int depth)
	{
		Vec3 rs;
		result.w = 1.0f;
		RenderThreadStorage * storage = (RenderThreadStorage *)TlsGetValue(scene->GetSystem()->TlsId);
		scene->GetGlobalPhotonMap()->EstimateIrradiance(rs, storage->PhotonSearchBuffer, dgeom->Position, dgeom->Normal, scene->GetGlobalPhotonMapRadius());
		result.x = rs.x;
		result.y = rs.y;
		result.z = rs.z;
	}

}